#!/usr/bin/env python
# coding: utf-8

from dashscope.api_entities.dashscope_response import Role
from flask import Flask, request, jsonify
from flask_cors import CORS
import requests


# 封装模型响应函数（添加超时控制）
def get_response(messages, model="deepseek-r1", timeout=10):
    # 将消息数组转换为字符串格式
    prompt_str = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
    print(f"prompt_str: {prompt_str}")
    url = "http://localhost:11434/api/generate"
    data = {
        "model": model,
        "prompt": prompt_str,
        "stream": True,
    }

    try:
        # 添加超时参数
        response = requests.post(url, json=data, timeout=timeout)
        print(f"大模型返回数据：{response.json()}")
        if response.status_code == 200:
            return response.json()
        else:
            return None
    except requests.exceptions.Timeout:
        print(f"请求超时：{timeout}秒内未获得响应")
        return None
    except requests.exceptions.RequestException as e:
        print(f"请求异常: {str(e)}")
        return None


app = Flask(__name__)
CORS(app)


# 读取 system prompt 文件内容
def load_system_prompt():
    with open("sec0721/system_prompt_local.md", "r", encoding="utf-8") as f:
        return f.read()


@app.route("/ai/deepseek/secretary/get-response", methods=["POST"])
def api_get_response():
    data = request.get_json()
    messages = data.get("messages")
    review = data.get("review")

    if messages:
        # 多轮对话，直接用传入的 messages
        pass
    elif review:
        # 单轮对话，兼容原有逻辑
        system_prompt = load_system_prompt()  # 读取专业 prompt
        messages = [
            {
                "role": "system",
                "content": system_prompt,
            },
            {"role": "user", "content": review},
        ]
    else:
        return jsonify({"error": "Missing review or messages parameter"}), 400

    response_data = get_response(messages)
    if not response_data:
        return (
            jsonify(
                {
                    "code": 500,
                    "error": "Failed to get response from model",
                    "data": None,
                }
            ),
            500,
        )

    # 从解析后的JSON中获取模型输出
    content = response_data.get("response", "")
    return jsonify({"code": 0, "data": content})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)
